########################################################################
#    Classes and methods for dlmm
# 
#    Author: Wayne Zhang, June 2012 
#            actuary_zhang@hotmail.com
########################################################################




#------------------------------ Classes -------------------------------#

setOldClass("bugs")

#' The \code{dlmm} class
#' 
#' This is the class of the object returned from \code{dlmm}. 
#' 
#' 
#' @section Slots:
#'  \describe{
#' \item{call}{the matched call.}
#' \item{data}{a list containing the response vector (\code{Y}), the design matrices 
#'  for the fixed effects (\code{X}), the random effects (\code{Z1}, \code{Z2}, ...) 
#'  and the dynamic effects (\code{F}), the grouping factor for the random effects
#'  (\code{grp1}, \code{grp2}, ...) and the time factor (\code{time}). Some of these
#'  elements may not be present if the correponding effects are not specified.}
#' \item{inits}{a list of initial values}
#' \item{parameters.to.save}{the names of the parameters to be stored, an object of 
#'  class \code{"character"}.}
#' \item{dims}{a vector of dimensions for the design matrices, an object of class 
#'  \code{"integer"}.}
#' \item{sims}{the simulation result, an object of class \code{"bugs"}. See 
#' \code{\link[R2WinBUGS]{bugs}} for detailed definition of this class.}
#' }
#' 
#' @docType class
#' @rdname class
#' @exportClass dlmm
#' @aliases dlmm-class
#' @name dlmm-class
#' 
setClass("dlmm",
  representation(
    call = "call",
    data = "list", 
    inits = "list", 
    parameters.to.save = "character", 
    dims = "integer", 
    sims = "bugs"))


#------------------------------ Methods -------------------------------#
#' Methods defined for the \code{dlmm} class
#'
#' \bold{\code{dynef}}: generic function for extracting dynamic effects
#' 
#' @param object an object of class \code{"dlmm"}
#'
#' @docType methods
#' @rdname methods
#' @name dlmm-methods
#' @aliases dynef
#' 
#' 
setGeneric("dynef",
  function(object, ...)
    standardGeneric("dynef")
)



#' 
#' @param object an object of class \code{"dlmm"}
#' @param type a character string indicating whether the median or 
#'  the mean should be used 
#' @param sd whether standard deviations should be computed. If so, 
#'  they are stored in the \code{"sd"} attribute of the returned result.
#' @param quantile a numeric vector of probabilities with values in \code{[0,1]}.
#'  These are used to compute the empirical quantiles. 
#' @param ... not used
#'   
#' @docType methods
#' @rdname methods
#' @aliases dynef,dlmm-method
#' @exportMethod dynef
#' @name dlmm-methods
#' 
setMethod("dynef", signature = "dlmm",
  function(object, type = c("median", "mean"), sd = FALSE, 
           quantile = NULL, ...){
    type <- match.arg(type)
    cn <- ifelse(type == "median", 5, 1)
    s <- object@sims$summary
    rn <- grep("^theta\\[", rownames(s))
    dm <- object@dims
    # point estimates
    mu.theta <- matrix(s[rn, cn], dm["n.t"], dm["n.theta"], byrow = TRUE)
    colnames(mu.theta) <- colnames(object@data$F)
    # standard errors
    if (sd){
      sd.theta <- matrix(s[rn, 2], dm["n.t"], dm["n.theta"], byrow = TRUE)
      attr(mu.theta, "sd") <- sd.theta
    }
    # quantiles
    if (!is.null(quantile)){
      qt.theta <- lapply(1:dm["n.theta"], function(x)
        apply(object@sims$sims.list$theta[ , , x], 2, quantile, quantile))
      qt.theta <- lapply(qt.theta, t)
      attr(mu.theta, "quantile") <- qt.theta
    }
    return(mu.theta)
  }
)

#' 
#' \bold{\code{fixef}}: extract fixed effects
#' 
#' @inheritParams dlmm-methods
#' @importMethodsFrom lme4 fixef ranef
#' @rdname methods
#' @aliases fixef,dlmm-method
#' @exportMethod fixef
#' 
setMethod("fixef", signature = "dlmm",
  function(object, type = c("median", "mean"), sd = FALSE,
           quantile = NULL, ...){
    if (object@dims["n.beta"]){
      type <- match.arg(type)
      cn <- ifelse(type == "median", 5, 1)
      s <- object@sims$summary
      rn <- grep("^beta\\[", rownames(s))
      dm <- object@dims
      mu.beta <- as.numeric(s[rn, cn])
      names(mu.beta) <- colnames(object@data$X)
      if (sd){
        sd.beta <- as.numeric(s[rn, 2])
        attr(mu.beta, "sd") <- sd.beta
      }
      sim.beta <- object@sims$sims.matrix[, rn, drop = FALSE]
      if (!is.null(quantile)){
        qt.beta <- t(apply(sim.beta, 2, quantile, quantile))
        attr(mu.beta, "quantile") <- qt.beta
      }
      return(mu.beta)
    } else {
      return(NULL)
    }
  }
)

#' \bold{\code{VarCorr}}: extract the variance component estimation.
#' 
#' @importMethodsFrom lme4 VarCorr
#' @rdname methods
#' @aliases VarCorr,dlmm-method
#' @exportMethod VarCorr
#'
setMethod("VarCorr", signature(x = "dlmm"),
  function(x, ...){
    s <- x@sims$summary
    # random effects variance 
    zn <- grep("^Z\\d+$", names(x@data))
    nc <- unname(sapply(x@data[zn], ncol))
    nm <- lapply(x@data[zn], colnames)
    rn <- grep("^sigma.b\\d+", rownames(s))
    vc <- s[rn, 5]
    vcl <-  vector("list", length(nc))
    Gp <- c(0, cumsum(nc * nc))
    for (i in 1:length(nc)){
      ln <- vc[(Gp[i] + 1):Gp[i + 1]]
      vcl[[i]] <- matrix(ln, nc[i])
      dimnames(vcl[[i]]) <- list(nm[[i]], nm[[i]])
      if (nc[i] == 1)
        vcl[[i]] <- vcl[[i]]^2    
    }
    # dynamic effects variance
    ps <- grep("^sigma.theta", rownames(s))
    vt <- matrix(unname(s[ps, 5]), x@dims["n.theta"])
    if (x@dims["n.theta"] == 1)
      vt <- vt^2
    dimnames(vt) <- rlply(2, colnames(x@data$F))
    vcl$vt <- vt
    # get correlations
    ans <- lapply(vcl, function(v) {
      stddev <- sqrt(diag(v))
      correl <- t(v/stddev)/stddev
      diag(correl) <- 1
      attr(v, "stddev") <- stddev
      attr(v, "correlation") <- correl
      v
    })
    
    # assign names
    names(ans) <- factor_name(x, TRUE)    
    ps <- grep("sigma.y$", rownames(s))
    attr(ans, "sc") <- ifelse(unname(x@dims["is.gauss"]), s[ps, 5], NA) 
    ans
  })

#' \code{summary}: produce summary report for the fitted model.
#' 
#' @importMethodsFrom lme4 summary
#' @rdname methods
#' @aliases summary,dlmm-method
#' @exportMethod summary
#'
setMethod("summary", signature = "dlmm",
  function(object, ...)
    print_dlmm(object)
)


#' \code{show}: the default print (show) method for a dlmm object.
#' 
#' @importMethodsFrom methods show
#' @rdname methods
#' @aliases show,dlmm-method
#' @exportMethod show
#'
setMethod("show", signature = "dlmm",  
  function(object) 
    print_dlmm(object)
)


# get the name of the grouping factors 
factor_name <- function(x, prefix = FALSE){
  ps <- grep("(Z\\d)|(^F$)", names(x@data))
  nm <- unname(sapply(x@data[ps], attr, "factor"))
  if (prefix){
    ll <- length(nm)
    if (x@dims["n.term"]){
      c(paste0("random:", nm[1:(ll - 1)]),
        paste0("dynamic:", nm[ll]))
    } else{
      paste0("dynamic:", nm[ll])
    }
  } else{
    nm
  }
}

# get # levels for each grouping factor
factor_nlev <- function(x){
  ps <- grep("(grp\\d)|(^time$)", names(x@data))
  nlev <- sapply(x@data[ps], function(x) length(unique(x)))
  names(nlev) <- factor_name(x)
  nlev
}


# print out (summarize) model results
print_dlmm <- function(x, digits = max(3, getOption("digits") - 3)){
  sims <- x@sims
  # fixed effects
  fcoef <- fixef(x, sd = TRUE, pvalue = TRUE, quantile = c(0.025, 0.975))
  coefs <- cbind("Estimate" = fcoef, "Std. Error" = attr(fcoef, "sd"),
                 "Lower (2.5%)" = attr(fcoef, "quantile")[, 1],
                 "Upper (97.5%)" = attr(fcoef, "quantile")[, 2])
  
  # variance components  
  varcor <- VarCorr(x)
  REmat <- lme4:::formatVC(varcor)
  if (is.na(attr(varcor, "sc")))
    REmat <- REmat[-nrow(REmat), , drop = FALSE]  
  
  # start printing
  cat("Dynamic linear mixed model fitted via Bugs")
  if (!is.null(sims$model.file)) 
    cat(" (model at \"", sims$model.file, "\")\n", sep = "")
  cat(sims$n.chains, " chains, each with ", sims$n.iter, " iterations (first ", 
      sims$n.burnin, " discarded)", sep = "")
  if (sims$n.thin > 1) 
    cat(", n.thin =", sims$n.thin)
  cat("\n n.sims =", sims$n.sims, "iterations saved\n")
  cat("\n")
  if (!is.null(x@call$formula))
    cat("Formula:", deparse(x@call$formula),"\n")
  if (!is.null(x@call$data))
    cat("   Data:", deparse(x@call$data), "\n")
  if (!is.null(x@call$subset))
    cat(" Subset:", deparse(x@call$subset),"\n")
  
  cat("\nRandom and dynamic variance components:\n")
  print(REmat, quote = FALSE, digits = digits)
  
  ngrps <- factor_nlev(x)
  cat(sprintf("Number of obs: %d, groups: ", x@dims["n.obs"]))
  cat(paste(paste(names(ngrps), ngrps, sep = ", "), collapse = "; "))
  cat("\n")
  
  if (nrow(coefs) > 0) {
    cat("\nFixed effects:\n")
    printCoefmat(coefs, zap.ind = 3, digits = digits)
    cat("\n")    
  }
  
  pD <- sims$pD
  DIC <- sims$DIC
  DICframe <- with(sims, data.frame(DIC = DIC, pD = pD,
                                    row.names = ""))
  if (sims$isDIC) {
    cat(paste("DIC info: ", sep = ""))
    if (length(sims$DIC) == 1) {
      cat("pD =", R2WinBUGS:::fround(sims$pD, 1), "and DIC =", 
          R2WinBUGS:::fround(sims$DIC, 1))
    }
    else if (length(sims$DIC) > 1) {
      print(round(sims$DIC, 1))
    }
    cat("\nDIC is an estimate of expected predictive error (lower deviance is better).\n")
  }
  
  out <- list(fixef = coefs, VarCorr = REmat, DIC = DICframe, ngrps = ngrps)
  invisible(out)
}


#' \bold{\code{refit}}: refit a \code{dlmm} object with modified structure. This is 
#' useful when the desired model is not a standard one supported by the current 
#' functionality. The user can use the \code{dlmm_bugs} function with argument 
#' \code{do.fit = FALSE} to set up the basic model structure. The user then needs
#' to modify the bugs script stored in \code{model.file}, the underlying \code{data}, 
#' or the initial values \code{inits} to accommodate the desired model. A call of 
#' \code{refit} then fits the model with the new setup. 
#' 
#' @param newrep not used 
#' @importMethodsFrom lme4 refit
#' @rdname methods
#' @aliases refit,dlmm,missing-method
#' @exportMethod refit
#'
setMethod("refit", signature(object = "dlmm", newresp = "missing"),  
  function(object, newresp, ...){
    mc <- getCall(object)
    # remove arguments not used in bugs
    m <- !(names(mc) %in% c("formula", "family", "data", "parameters.to.save", 
                            "weights", "offset", "subset", "na.action", "contrasts", 
                            "do.fit", "save.ranef"))
    mcl <- mc[m]
    # construct an unevaluated bugs call
    mcl <- as.call(c(as.list(mcl), list(data = object@data, inits = object@inits, 
                      parameters.to.save = object@parameters.to.save)))
    mcl[[1]] <- as.name("bugs")
    object@sims <- eval(mcl)
    object
  }    
)

